"""
NPM 包解压器

这个脚本用于自动化处理大量 NPM 包（.tgz 或 .tar.gz 文件）的解压过程。
它会递归地在源目录中查找包文件，源目录结构应包含 'benign' 和 'malicious' 子目录，
然后将所有文件解压到对应的目标目录结构中。

主要功能:
- 解压所有文件类型（不再限制为仅 .js 和 package.json）
- 处理 NPM tar 包中常见的 'package/' 目录前缀
- 提供路径遍历攻击防护
- 在目标位置重建源目录结构
- 报告进度并汇总成功和失败情况

依赖:
- tqdm: 用于用户友好的进度条显示。使用 `pip install tqdm` 安装。

使用方法:
    python extract_packages.py
    
    （请修改文件底部主执行块中的 SOURCE_ROOT 和 TARGET_ROOT 路径以匹配你的环境）
"""
import os
import sys
import tarfile
import shutil
from pathlib import Path
from tqdm import tqdm
from typing import Optional

# 在 Windows 上设置控制台编码为 UTF-8 以正确显示字符
if sys.platform == "win32":
    sys.stdout.reconfigure(encoding='utf-8')
    os.system('chcp 65001 >nul')

def npm_package_filter(member: tarfile.TarInfo, path: str = None) -> Optional[tarfile.TarInfo]:
    """
    NPM 包解压的自定义过滤器，用于 tarfile.extractall 选择性解压文件。

    此过滤器执行以下操作:
    1. 安全检查: 防止路径遍历攻击（如 '..'、绝对路径）
    2. 路径标准化: 移除 NPM tar 包中常见的 'package/' 前缀
    3. 文件选择: 允许所有文件类型被解压
    4. 目录创建: 允许创建目录以维持结构

    参数:
        member: tar 归档中成员的 TarInfo 对象
        path: 目标路径（Python 3.12+ 会传递此参数，为兼容性保留）

    返回:
        如果成员应该被解压则返回修改后的 TarInfo 对象，否则返回 None
    """
    # 1. 安全检查
    if '..' in member.name or member.name.startswith('/'):
        print(f"  [跳过] 检测到不安全路径: {member.name}")
        return None

    # 2. 路径标准化
    # NPM tar 包通常有一个 'package/' 根目录，我们需要移除它
    if member.name.startswith('package/'):
        member.name = member.name[len('package/'):]
        # 如果名称现在为空（即它是 'package/' 目录本身），跳过它
        if not member.name:
            return None

    # 3. 文件选择 & 4. 目录创建
    if member.isdir():
        return member  # 允许创建目录

    if member.isfile():
        return member  # 允许所有文件类型

    # 拒绝其他类型（如符号链接等特殊文件）
    return None

def extract_packages(source_dir: Path, target_dir: Path):
    """
    从源目录查找并解压所有 NPM 包到目标目录。

    直接处理源目录中的所有 .tgz 和 .tar.gz 文件，不再假设特定的子目录结构。
    适用于处理真实世界收集的未标记数据。

    参数:
        source_dir: 包含包归档的目录（可以是任意结构）
        target_dir: 解压后文件将被放置的根目录
                   结果结构: target_dir/package_name/
    """
    print(f"源目录: {source_dir}")
    print(f"目标目录: {target_dir}")
    print(f"Python 版本: {sys.version}")

    # 创建基础目标目录
    target_dir.mkdir(parents=True, exist_ok=True)
    
    extracted_count = 0  # 成功解压计数
    failed_count = 0     # 失败解压计数
    failed_packages = [] # 失败包列表，用于详细报告
    
    # 递归查找所有 .tgz 和 .tar.gz 文件
    print(f"\n正在扫描目录查找包文件...")
    tgz_files = []
    
    # 使用 rglob 递归查找所有 .tgz 和 .tar.gz 文件
    tgz_files.extend(list(source_dir.rglob('*.tgz')))
    tgz_files.extend(list(source_dir.rglob('*.tar.gz')))
    
    print(f"找到 {len(tgz_files)} 个包文件。")
    
    if not tgz_files:
        print(f"在目录 {source_dir} 中没有找到任何 .tgz 或 .tar.gz 文件")
        return

    # 使用 tqdm 显示进度条
    for tgz_file in tqdm(tgz_files, desc="解压包", unit="包"):
        # 生成包名（移除 .tgz 或 .tar.gz 扩展名）
        package_name = tgz_file.name.replace('.tgz', '').replace('.tar.gz', '')
        
        # 如果包名包含路径分隔符，替换为下划线以避免目录冲突
        safe_package_name = package_name.replace('/', '_').replace('\\', '_')
        
        extract_path = target_dir / safe_package_name
        
        try:
            # 如果目标目录已存在，先删除以确保干净的解压
            if extract_path.exists():
                shutil.rmtree(extract_path)
            
            # 创建目标目录
            extract_path.mkdir(parents=True, exist_ok=True)
            
            # 使用自定义过滤器打开和解压归档
            with tarfile.open(tgz_file, 'r:gz') as tar:
                # 检查 Python 版本以决定使用哪种方法
                python_version = sys.version_info
                if python_version >= (3, 12):
                    # Python 3.12+ 使用新的 filter 参数
                    tar.extractall(path=extract_path, filter=npm_package_filter)
                else:
                    # 旧版本 Python 使用 members 参数
                    filtered_members = []
                    for member in tar.getmembers():
                        filtered_member = npm_package_filter(member)
                        if filtered_member is not None:
                            filtered_members.append(filtered_member)
                    tar.extractall(path=extract_path, members=filtered_members)
            
            extracted_count += 1
            
        except (tarfile.TarError, EOFError, PermissionError, OSError) as e:
            failed_count += 1
            error_reason = str(e)
            failed_packages.append((tgz_file.name, error_reason))
            
            # 清理部分解压/失败的目录
            if extract_path.exists():
                try:
                    shutil.rmtree(extract_path)
                except Exception:
                    pass  # 忽略清理错误
                    
        except Exception as e:
            # 捕获其他未预期的异常
            failed_count += 1
            error_reason = f"未知错误: {str(e)}"
            failed_packages.append((tgz_file.name, error_reason))
            
            # 清理部分解压/失败的目录
            if extract_path.exists():
                try:
                    shutil.rmtree(extract_path)
                except Exception:
                    pass  # 忽略清理错误
    
    # --- 最终摘要 ---
    print(f"\n解压完成:")
    print(f"  成功解压: {extracted_count} 个包")
    print(f"  解压失败: {failed_count} 个包")
    print("注意: 所有文件类型都已从每个包中解压出来")
    
    # 如果有失败的包，显示详细信息
    if failed_count > 0:
        print(f"\n--- 失败详情 ---")
        for pkg_name, reason in failed_packages:
            print(f"  - {pkg_name}: {reason}")
    
    # 显示解压结果的目录结构信息
    print(f"\n--- 解压结果 ---")
    if target_dir.exists():
        subdirs = [d for d in target_dir.iterdir() if d.is_dir()]
        print(f"  总共解压了 {len(subdirs)} 个包目录到: {target_dir}")
        
        # 显示前几个解压的包作为示例
        if subdirs:
            print(f"  示例包目录:")
            for subdir in subdirs[:5]:  # 只显示前5个
                print(f"    - {subdir.name}")
            if len(subdirs) > 5:
                print(f"    ... 以及其他 {len(subdirs) - 5} 个包")
def extract_packages_simple(source_dir: str, target_dir: str):
    """
    简化版的包解压函数，用于程序调用
    
    参数:
        source_dir: 源目录路径（字符串）
        target_dir: 目标目录路径（字符串）
    
    返回:
        tuple: (成功数量, 失败数量)
    """
    source_path = Path(source_dir)
    target_path = Path(target_dir)
    
    if not source_path.exists():
        raise ValueError(f"源目录不存在: {source_path}")
    
    # 创建目标目录
    target_path.mkdir(parents=True, exist_ok=True)
    
    # 查找所有包文件
    tgz_files = list(source_path.rglob('*.tgz')) + list(source_path.rglob('*.tar.gz'))
    
    extracted_count = 0
    failed_count = 0
    
    for tgz_file in tqdm(tgz_files, desc="解压", disable=False):
        package_name = tgz_file.name.replace('.tgz', '').replace('.tar.gz', '')
        safe_package_name = package_name.replace('/', '_').replace('\\', '_')
        extract_path = target_path / safe_package_name
        
        try:
            if extract_path.exists():
                shutil.rmtree(extract_path)
            
            extract_path.mkdir(parents=True, exist_ok=True)
            
            with tarfile.open(tgz_file, 'r:gz') as tar:
                python_version = sys.version_info
                if python_version >= (3, 12):
                    tar.extractall(path=extract_path, filter=npm_package_filter)
                else:
                    filtered_members = []
                    for member in tar.getmembers():
                        filtered_member = npm_package_filter(member)
                        if filtered_member is not None:
                            filtered_members.append(filtered_member)
                    tar.extractall(path=extract_path, members=filtered_members)
            
            extracted_count += 1
            
        except Exception:
            failed_count += 1
            if extract_path.exists():
                try:
                    shutil.rmtree(extract_path)
                except Exception:
                    pass
    
    return extracted_count, failed_count

def main():
    """主函数，配置路径并执行解压操作"""
    import argparse
    
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(
        description="NPM包解压器 - 解压指定目录下的所有.tgz文件",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  python extract_packages.py --source sentinel_data/packages/2025-08-11 --target extracted_packages
  python extract_packages.py -s sentinel_data/test -t analysis_runs/extracted
        """
    )
    
    parser.add_argument(
        '--source', '-s',
        type=str,
        required=True,
        help='包含.tgz文件的源目录路径'
    )
    
    parser.add_argument(
        '--target', '-t', 
        type=str,
        required=True,
        help='解压后文件的目标目录路径'
    )
    
    parser.add_argument(
        '--verbose', '-v',
        action='store_true',
        help='显示详细输出'
    )
    
    args = parser.parse_args()
    
    # 转换为 Path 对象
    source_path = Path(args.source)
    target_path = Path(args.target)
    
    # 验证源目录是否存在
    if not source_path.exists():
        print(f"错误: 源目录不存在: {source_path}")
        print("请检查路径是否正确")
        return 1
    
    if not source_path.is_dir():
        print(f"错误: 源路径不是一个目录: {source_path}")
        return 1
    
    print("=" * 60)
    print("NPM 包解压器 (通用版)")
    print("=" * 60)
    print(f"源目录: {source_path.absolute()}")
    print(f"目标目录: {target_path.absolute()}")
    
    if args.verbose:
        print(f"详细模式: 启用")
    
    # 询问用户确认
    try:
        confirm = input("\n是否继续解压? [y/N]: ").strip().lower()
        if confirm not in ['y', 'yes']:
            print("用户取消操作")
            return 0
    except KeyboardInterrupt:
        print("\n用户中断操作")
        return 0
    
    # --- 执行解压 ---
    try:
        extract_packages(source_dir=source_path, target_dir=target_path)
        print(f"\n✅ 解压操作已完成。请检查目标目录: {target_path}")
        return 0
    except KeyboardInterrupt:
        print("\n\n用户中断了解压过程")
        return 1
    except Exception as e:
        print(f"\n❌ 解压过程中发生错误: {e}")
        if args.verbose:
            import traceback
            traceback.print_exc()
        return 1
if __name__ == "__main__":
    main()